// SVMClassificationModule.h

#pragma once
#include <map>
#include "SVM.h"
#include "FeatureVectors.h"
#include "NormalizationModule.h"
#include "FeatureExtractorBase.h"

class CSVMClassificationModule
{
	public:
		CSVMClassificationModule();
		virtual ~CSVMClassificationModule();
		/**
			Loads/Saves an SVM, there are three files
			1. The actual SVM
			2. The normalization constants
			3. The class labels files. This translates between class names and double values.
		*/
		bool Load( std::string fname, std::string path );
		bool Save( std::string fname, std::string path );

		/**
			TClassToNameMap is a mapping between double values and class
			names. E.g. 1 = cars, -1 = elephants
		*/
		void SetClassLabels( TClassToNameMap& classLabels );
		TClassToNameMap GetClassLabels( void );

		/**
			Get the class name from a double value; 
		*/
		std::string GetClassName( double value ); 

		/**
			Set/Get a feature extractor. A feature extractor takes in an image and returns a feature vector
		*/
		void SetFeatureExtractor( CFeatureExtractorBase * extractor ); 
		CFeatureExtractorBase* GetFeatureExtractor( void );  

		/**
			set the libsvm parameters (kernel type, etc).
		*/
		void SetParams( svm_parameter params ); 
		svm_parameter GetParams( void ); 

		/**
			Build the SVM using a single datapoint manager with class labels assigned
			if manageImgs is true then the system loads the images, and then unloads them. 
		*/
		bool BuildSVM( CDataPointManager * data, bool manageImgs = true ); 
		/**
			Alternatively build the SVM using a datapoint manager for each with class labels in the classes vectors
			if manageImgs is true then the system loads the images, and then unloads them. 
		*/
		bool BuildSVM( TDataMgrVector& data, TTruthLabels classes, bool manageImgs = true);


		/**
			Takes in an image, returns a classification
		*/
		double Classify( IplImage* image );

		/** 
			Classify a single image, set the datapoint's experimental classification value.
			Also return the classification as a single double.
		*/
		double Classify( CTruthDataPoint* datapoint, bool manageImgs = true);

		/*
			Classify a whole dataset, set the classification value for each datapoint,
			also return the class labels. 
		*/
		TTruthLabels Classify( CDataPointManager * data, bool manageImgs = true); 

		std::string GetTheClassName( double value ){return m_ClassNameMap[value];};
protected:
		bool LoadSVMFile(std::string fname, std::string path);
		bool LoadClassNameFile( std::string fname, std::string path ); 
		bool SaveSVMFile(std::string fname, std::string path);
		bool SaveClassNameFile( std::string fname, std::string path ); 
		std::string m_classNameEXT;
		std::string m_svmNameEXT;
		CFeatureExtractorBase *m_extractor;
		CNormalizationModule m_normalizer; 
		svm_parameter m_parameters; 
		svm_model *m_svmModel; 
		svm_problem m_problem; 
		TClassToNameMap m_ClassNameMap;

};

